import torch
from ultralytics import YOLO
import numpy as np
from PIL import Image


class YOLOProposer:
    """
    Stage 1: YOLO-based Region Proposal Network

    Supports tile-based processing for large images.
    """

    def __init__(self, weights_path=None, device='cpu'):
        print("Initializing Stage 1: YOLO Proposal Network...")
        self.model = YOLO(weights_path)
        self.model.to(device)
        self.device = device
        print(f"YOLO Proposal Network loaded with weights from '{weights_path}'.")

    @torch.no_grad()
    def propose(self, image_path, tile_size=None, tile_overlap=None,
                conf_thresh=0.1, iou_thresh=0.5):
        """
        Generate region proposals (ROIs) from input image with optional tiling.

        Args:
            image_path: Path to input image
            tile_size: Tile size in pixels; None for full image inference
            tile_overlap: Overlap between adjacent tiles in pixels
            conf_thresh: Confidence threshold for proposal filtering
            iou_thresh: IoU threshold for Non-Maximum Suppression

        Returns:
            proposals: Numpy array of shape [N, 4] with format [x1, y1, x2, y2]
        """
        # Direct inference for small images or when tiling is disabled
        if tile_size is None:
            return self._predict_full_image(image_path, conf_thresh, iou_thresh)

        # Tile-based inference for large images
        return self._predict_with_tiles(image_path, tile_size, tile_overlap, conf_thresh, iou_thresh)

    def _predict_full_image(self, image_path, conf_thresh, iou_thresh):
        """Perform inference on the full image without tiling."""
        results = self.model.predict(
            source=image_path,
            conf=conf_thresh,
            iou=iou_thresh,
            verbose=False,
            imgsz=1280,
            max_det=300
        )

        if len(results) == 0 or len(results[0].boxes) == 0:
            return np.empty((0, 4))

        boxes = results[0].boxes.xyxy.cpu().numpy()
        return boxes

    def _predict_with_tiles(self, image_path, tile_size, tile_overlap, conf_thresh, iou_thresh):
        """Process large images using tile-based inference strategy."""
        img = Image.open(image_path).convert("RGB")
        img_w, img_h = img.size

        # Fall back to full image inference for small images
        if img_w <= tile_size and img_h <= tile_size:
            return self._predict_full_image(image_path, conf_thresh, iou_thresh)

        all_boxes = []
        stride = tile_size - tile_overlap

        # Generate tile grid
        for y in range(0, img_h, stride):
            for x in range(0, img_w, stride):
                # Compute tile boundaries
                x1 = x
                y1 = y
                x2 = min(x + tile_size, img_w)
                y2 = min(y + tile_size, img_h)

                # Extract tile region
                tile = img.crop((x1, y1, x2, y2))

                # Save tile to temporary file
                import tempfile
                with tempfile.NamedTemporaryFile(suffix='.jpg', delete=False) as tmp:
                    tile.save(tmp.name)
                    tile_path = tmp.name

                # Run inference on tile
                results = self.model.predict(
                    source=tile_path,
                    conf=conf_thresh,
                    iou=iou_thresh,
                    verbose=False,
                    imgsz=tile_size,
                    max_det=100
                )

                # Clean up temporary file
                import os
                os.unlink(tile_path)

                if len(results) > 0 and len(results[0].boxes) > 0:
                    # Extract boxes and transform to global coordinates
                    tile_boxes = results[0].boxes.xyxy.cpu().numpy()
                    # Apply spatial offset to global coordinate system
                    tile_boxes[:, [0, 2]] += x1
                    tile_boxes[:, [1, 3]] += y1
                    all_boxes.append(tile_boxes)

        if len(all_boxes) == 0:
            return np.empty((0, 4))

        # Merge all tile predictions
        all_boxes = np.vstack(all_boxes)

        # Apply global NMS to remove duplicate detections across tiles
        from src.utils.postprocess import non_max_suppression_global
        scores = np.ones(len(all_boxes))  # YOLO pre-filtered by confidence, use 1.0
        all_boxes = non_max_suppression_global(all_boxes, scores, iou_thresh)

        return all_boxes